import cv2
import numpy as np
from PIL import Image
from io import BytesIO
import logging

logger = logging.getLogger(__name__)

def enhance_image(image_data: bytes) -> bytes:
    """Enhance image quality for better OCR accuracy
    
    Performs the following operations:
    1. Converts to grayscale
    2. Applies median filtering to remove noise
    3. Applies CLAHE (Contrast Limited Adaptive Histogram Equalization) for contrast adjustment
    4. Applies sharpening filter
    
    Args:
        image_data: Original image data as bytes
        
    Returns:
        Enhanced image data as bytes
    """
    try:
        # Load image from bytes
        pil_image = Image.open(BytesIO(image_data))
        
        # Convert PIL image to numpy array (RGB)
        img_array = np.array(pil_image)
        
        # Convert RGB to BGR for OpenCV
        img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        
        # Convert to grayscale
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        
        # Apply median filtering to remove noise
        denoised = cv2.medianBlur(gray, 3)
        
        # Apply CLAHE for contrast enhancement
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(denoised)
        
        # Apply sharpening filter
        kernel = np.array([[-1, -1, -1],
                          [-1,  9, -1],
                          [-1, -1, -1]])
        sharpened = cv2.filter2D(enhanced, -1, kernel)
        
        # Convert back to PIL Image
        enhanced_pil = Image.fromarray(sharpened)
        
        # Save to bytes
        img_byte_arr = BytesIO()
        enhanced_pil.save(img_byte_arr, format=pil_image.format or 'PNG')
        enhanced_bytes = img_byte_arr.getvalue()
        
        logger.info("Image enhancement completed successfully")
        return enhanced_bytes
        
    except Exception as e:
        logger.error(f"Image enhancement failed: {e}")
        # Return original image if enhancement fails
        return image_data

def correct_skew(image_data: bytes) -> bytes:
    """Correct image skew using Hough transform
    
    Args:
        image_data: Original image data as bytes
        
    Returns:
        Skew-corrected image data as bytes
    """
    try:
        # Load image from bytes
        pil_image = Image.open(BytesIO(image_data))
        
        # Convert PIL image to numpy array (RGB)
        img_array = np.array(pil_image)
        
        # Convert RGB to BGR for OpenCV
        img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        
        # Convert to grayscale
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        
        # Apply threshold to get binary image
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        # Detect edges
        edges = cv2.Canny(binary, 50, 150, apertureSize=3)
        
        # Detect lines using Hough transform
        lines = cv2.HoughLines(edges, 1, np.pi / 180, threshold=100)
        
        if lines is not None:
            angles = []
            for rho, theta in lines[:, 0]:
                # Convert to degrees
                angle = np.degrees(theta)
                # Normalize angle to -45 to 45 degrees
                if angle < 45:
                    angles.append(angle)
                elif angle > 135:
                    angles.append(angle - 180)
            
            if angles:
                # Calculate median angle
                median_angle = np.median(angles)
                
                # Only correct if skew is significant (> 1 degree)
                if abs(median_angle) > 1.0:
                    # Get image dimensions
                    (h, w) = img_bgr.shape[:2]
                    center = (w // 2, h // 2)
                    
                    # Calculate rotation matrix
                    rotation_matrix = cv2.getRotationMatrix2D(center, median_angle, 1.0)
                    
                    # Perform rotation
                    corrected = cv2.warpAffine(img_bgr, rotation_matrix, (w, h), 
                                             flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
                    
                    # Convert back to PIL Image (BGR to RGB)
                    corrected_rgb = cv2.cvtColor(corrected, cv2.COLOR_BGR2RGB)
                    corrected_pil = Image.fromarray(corrected_rgb)
                    
                    # Save to bytes
                    img_byte_arr = BytesIO()
                    corrected_pil.save(img_byte_arr, format=pil_image.format or 'PNG')
                    corrected_bytes = img_byte_arr.getvalue()
                    
                    logger.info(f"Image skew correction completed: {median_angle:.2f} degrees")
                    return corrected_bytes
        
        # Return original image if no significant skew detected
        logger.info("No significant skew detected, returning original image")
        return image_data
        
    except Exception as e:
        logger.error(f"Skew correction failed: {e}")
        # Return original image if correction fails
        return image_data

def preprocess_image_for_ocr(image_data: bytes, enhance: bool = True, correct_skew_flag: bool = True) -> bytes:
    """Complete preprocessing pipeline for OCR
    
    Args:
        image_data: Original image data as bytes
        enhance: Whether to apply image enhancement
        correct_skew_flag: Whether to apply skew correction
        
    Returns:
        Preprocessed image data as bytes
    """
    processed_data = image_data
    
    # Apply image enhancement if requested
    if enhance:
        processed_data = enhance_image(processed_data)
    
    # Apply skew correction if requested
    if correct_skew_flag:
        processed_data = correct_skew(processed_data)
    
    return processed_data